{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Testing extract_relevance_info_from_docs_with_conversation Function\n",
    "\n",
    "This notebook demonstrates how to use and test the `extract_relevance_info_from_docs_with_conversation` function from the ByzerLLM library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m2024-08-28 23:44:14.074\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mbyzerllm.utils.connect_ray\u001b[0m:\u001b[36mconnect_cluster\u001b[0m:\u001b[36m48\u001b[0m - \u001b[1mJDK 21 will be used (/Users/allwefantasy/.auto-coder/jdk-21.0.2.jdk/Contents/Home)...\u001b[0m\n",
      "2024-08-28 23:44:14,113\tINFO worker.py:1564 -- Connecting to existing Ray cluster at address: 127.0.0.1:6379...\n",
      "2024-08-28 23:44:14,128\tINFO worker.py:1740 -- Connected to Ray cluster. View the dashboard at \u001b[1m\u001b[32m127.0.0.1:8265 \u001b[39m\u001b[22m\n"
     ]
    }
   ],
   "source": [
    "import byzerllm\n",
    "from byzerllm import ByzerLLM\n",
    "\n",
    "# Initialize ByzerLLM\n",
    "llm = ByzerLLM.from_default_model(\"deepseek_chat\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare Test Data\n",
    "\n",
    "Let's create some sample documents and a conversation history to test our function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sample documents\n",
    "documents = [\n",
    "    \"The capital of France is Paris. It is known for its beautiful architecture and cuisine.\",\n",
    "    \"Python is a popular programming language used in data science and web development.\",\n",
    "    \"The Great Wall of China is one of the most famous landmarks in the world.\"\n",
    "]\n",
    "\n",
    "# Sample conversation history\n",
    "conversations = [\n",
    "    {\"role\": \"user\", \"content\": \"Tell me about France.\"},\n",
    "    {\"role\": \"assistant\", \"content\": \"France is a country in Western Europe known for its rich history and culture.\"},\n",
    "    {\"role\": \"user\", \"content\": \"What's the capital of France?\"}\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define and Test the Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@byzerllm.prompt()\n",
    "def extract_relevance_info_from_docs_with_conversation(conversations, documents) -> str:\n",
    "    \"\"\"\n",
    "    使用以下文档和对话历史来提取相关信息。\n",
    "\n",
    "    文档：\n",
    "    {% for doc in documents %}\n",
    "    {{ doc }}\n",
    "    {% endfor %}\n",
    "\n",
    "    对话历史：\n",
    "    {% for msg in conversations %}\n",
    "    <{{ msg.role }}>: {{ msg.content }}\n",
    "    {% endfor %}\n",
    "\n",
    "    请根据提供的文档内容、用户对话历史以及最后一个问题，提取并总结文档中与问题相关的重要信息。\n",
    "    如果文档中没有相关信息，请回复\"该文档中没有与问题相关的信息\"。\n",
    "    提取的信息尽量保持和原文中的一样，并且只输出这些信息。\n",
    "    \"\"\"\n",
    "\n",
    "# Test the function\n",
    "result = extract_relevance_info_from_docs_with_conversation.with_llm(llm).run(conversations=conversations, documents=documents)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analyze the Results\n",
    "\n",
    "Let's analyze the output of our function to see if it correctly extracted the relevant information from the documents based on the conversation history."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Extracted relevant information:\")\n",
    "print(result)\n",
    "\n",
    "print(\"\\nDoes the extracted information answer the user's last question?\")\n",
    "print(\"User's last question:\", conversations[-1]['content'])\n",
    "print(\"Relevant information found:\", \"Yes\" if \"Paris\" in result else \"No\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test with Different Scenarios\n",
    "\n",
    "Let's test our function with different scenarios to ensure it works correctly in various situations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Scenario 1: Question about a topic not in the documents\n",
    "conversations_scenario1 = [\n",
    "    {\"role\": \"user\", \"content\": \"What's the capital of Germany?\"},\n",
    "]\n",
    "\n",
    "result_scenario1 = extract_relevance_info_from_docs_with_conversation.with_llm(llm).run(conversations=conversations_scenario1, documents=documents)\n",
    "print(\"Scenario 1 Result:\")\n",
    "print(result_scenario1)\n",
    "\n",
    "# Scenario 2: Question about a topic in the documents, but not related to France\n",
    "conversations_scenario2 = [\n",
    "    {\"role\": \"user\", \"content\": \"Tell me about Python programming.\"},\n",
    "]\n",
    "\n",
    "result_scenario2 = extract_relevance_info_from_docs_with_conversation.with_llm(llm).run(conversations=conversations_scenario2, documents=documents)\n",
    "print(\"\\nScenario 2 Result:\")\n",
    "print(result_scenario2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "In this notebook, we've tested the `extract_relevance_info_from_docs_with_conversation` function with various scenarios. The function should effectively extract relevant information from the provided documents based on the conversation history and the user's last question. \n",
    "\n",
    "By analyzing the results, we can see how well the function performs in different situations, such as when the relevant information is present in the documents or when it's not available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Dict\n",
    "from autocoder.common import AutoCoderArgs, SourceCode\n",
    "from byzerllm.utils.client.code_utils import extract_code\n",
    "import json\n",
    "from loguru import logger\n",
    "\n",
    "@byzerllm.prompt()\n",
    "def extract_relevance_range_from_docs_with_conversation(\n",
    "    conversations: List[Dict[str, str]], documents: List[str]\n",
    ") -> str:\n",
    "    \"\"\"\n",
    "    使用以下文档和对话历史来提取相关信息。\n",
    "\n",
    "    文档：\n",
    "    {% for doc in documents %}\n",
    "    {{ doc }}\n",
    "    {% endfor %}\n",
    "\n",
    "    对话历史：\n",
    "    {% for msg in conversations %}\n",
    "    <{{ msg.role }}>: {{ msg.content }}\n",
    "    {% endfor %}\n",
    "\n",
    "    请根据提供的文档内容、用户对话历史以及最后一个问题，从文档中提取与问题相关的一个或者多个重要信息。\n",
    "    每一块重要信息由 start_str 和 end_str 组成。\n",
    "    返回一个 JSON 数组，每个元素包含 \"start_str\" 和 \"end_str\"，分别表示重要信息的起始和结束字符串。\n",
    "    确保 start_str 和 end_str 在原文中都是唯一的，不会出现多次，并且不会重叠。\n",
    "    start_str 和 end_str 应该尽可能短，但要确保它们在原文中是唯一的。\n",
    "    \n",
    "\n",
    "    如果文档中没有相关重要信息，请返回空数组 []。\n",
    "\n",
    "    示例1：\n",
    "    文档：这是一个示例文档。大象是陆地上最大的动物之一。它们生活在非洲和亚洲。猫是常见的宠物，它们喜欢捕鼠。    \n",
    "    问题：大象生活在哪里？\n",
    "    返回：[{\"start_str\": \"大象是陆地\", \"end_str\": \"在非洲和亚洲。\"}]\n",
    "\n",
    "    示例2：\n",
    "    文档：太阳系有八大行星。地球是太阳系中第三颗行星，有海洋，有沙漠，温度适宜，昼夜温差小，是目前已知唯一有生命的星球。月球是地球唯一的天然卫星。\n",
    "    问题：地球的特点是什么？\n",
    "    返回：[{\"start_str\": \"地球是太阳系\", \"end_str\": \"生命的星球。\"}]\n",
    "\n",
    "    示例3：\n",
    "    文档：苹果是一种常见的水果。它富含维生素和膳食纤维。香蕉也是一种受欢迎的水果，含有大量钾元素。\n",
    "    问题：橙子的特点是什么？\n",
    "    返回：[]\n",
    "\n",
    "    请返回严格的 JSON 格式。不要有任何多余的文字或解释。\n",
    "    \"\"\"  \n",
    "\n",
    "# result = extract_relevance_range_from_docs_with_conversation.with_llm(llm).run(conversations=conversations, documents=documents)\n",
    "\n",
    "\n",
    "# {\"conversation\":conversations, \"doc\":[doc.source_code]}\n",
    "conversations = None\n",
    "documents = None\n",
    "with open(\"/tmp/rag.json\", \"r\") as f:\n",
    "    lines = f.read().split(\"\\n\")\n",
    "    for i, line in enumerate(lines):\n",
    "        if line:\n",
    "            v = json.loads(line)\n",
    "            conversations = v[\"conversation\"]\n",
    "            documents = v[\"doc\"]\n",
    "\n",
    "def process_range_doc(doc, max_retries=3):\n",
    "    for attempt in range(max_retries):\n",
    "        content = \"\"\n",
    "        try:\n",
    "            extracted_info = extract_relevance_range_from_docs_with_conversation.with_llm(\n",
    "                llm\n",
    "            ).run(\n",
    "                conversations, [doc.source_code]\n",
    "            )\n",
    "\n",
    "            print(extracted_info)                                                             \n",
    "            json_str = extract_code(extracted_info)[0][1]\n",
    "            json_objs = json.loads(json_str)                                    \n",
    "                                                \n",
    "            for json_obj in json_objs:\n",
    "                start_str = json_obj[\"start_str\"]\n",
    "                end_str = json_obj[\"end_str\"]\n",
    "                start_index = doc.source_code.index(start_str)\n",
    "                end_index = doc.source_code.index(end_str) + len(end_str)\n",
    "                content += doc.source_code[start_index:end_index] + \"\\n\"  \n",
    "                print(f\"{start_str} - {end_str} : {doc.source_code[start_index:end_index]}\")                                      \n",
    "            \n",
    "            return SourceCode(\n",
    "                module_name=doc.module_name, source_code=content.strip()\n",
    "            )\n",
    "        except Exception as e:\n",
    "            if attempt < max_retries - 1:\n",
    "                logger.warning(f\"Error processing doc {doc.module_name}, retrying... (Attempt {attempt + 1}) attempts: {str(e)}\")                                        \n",
    "            else:\n",
    "                logger.error(f\"Failed to process doc {doc.module_name} after {max_retries} attempts: {str(e)}\")\n",
    "                return SourceCode(\n",
    "                module_name=doc.module_name, source_code=content.strip()\n",
    "            )\n",
    "\n",
    "print(documents[0])\n",
    "m = process_range_doc(SourceCode(\n",
    "                module_name=\"test\", source_code=documents[0]\n",
    "            ))\n",
    "\n",
    "\n",
    "print(m.source_code)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
